import logging
from typing import Dict, List, Optional

from llama_index.core.base.llms.types import ChatMessage
from llama_index.core.settings import Settings
from llama_index.core.tools import FunctionTool
from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)

# Prompt based on https://github.com/e2b-dev/ai-artifacts
CODE_GENERATION_PROMPT = """You are a skilled software engineer. You do not make mistakes. Generate an artifact. You can install additional dependencies. You can use one of the following templates:

1. code-interpreter-multilang: "Runs code as a Jupyter notebook cell. Strong data analysis angle. Can use complex visualisation to explain results.". File: script.py. Dependencies installed: python, jupyter, numpy, pandas, matplotlib, seaborn, plotly. Port: none.

2. nextjs-developer: "A Next.js 13+ app that reloads automatically. Using the pages router.". File: pages/index.tsx. Dependencies installed: nextjs@14.2.5, typescript, @types/node, @types/react, @types/react-dom, postcss, tailwindcss, shadcn. Port: 3000.

3. vue-developer: "A Vue.js 3+ app that reloads automatically. Only when asked specifically for a Vue app.". File: app.vue. Dependencies installed: vue@latest, nuxt@3.13.0, tailwindcss. Port: 3000.

4. streamlit-developer: "A streamlit app that reloads automatically.". File: app.py. Dependencies installed: streamlit, pandas, numpy, matplotlib, request, seaborn, plotly. Port: 8501.

5. gradio-developer: "A gradio app. Gradio Blocks/Interface should be called demo.". File: app.py. Dependencies installed: gradio, pandas, numpy, matplotlib, request, seaborn, plotly. Port: 7860.

Make sure to use the correct syntax for the programming language you're using.
"""


class CodeArtifact(BaseModel):
    commentary: str = Field(
        ...,
        description="Describe what you're about to do and the steps you want to take for generating the artifact in great detail.",
    )
    template: str = Field(
        ..., description="Name of the template used to generate the artifact."
    )
    title: str = Field(..., description="Short title of the artifact. Max 3 words.")
    description: str = Field(
        ..., description="Short description of the artifact. Max 1 sentence."
    )
    additional_dependencies: List[str] = Field(
        ...,
        description="Additional dependencies required by the artifact. Do not include dependencies that are already included in the template.",
    )
    has_additional_dependencies: bool = Field(
        ...,
        description="Detect if additional dependencies that are not included in the template are required by the artifact.",
    )
    install_dependencies_command: str = Field(
        ...,
        description="Command to install additional dependencies required by the artifact.",
    )
    port: Optional[int] = Field(
        ...,
        description="Port number used by the resulted artifact. Null when no ports are exposed.",
    )
    file_path: str = Field(
        ..., description="Relative path to the file, including the file name."
    )
    code: str = Field(
        ...,
        description="Code generated by the artifact. Only runnable code is allowed.",
    )


class CodeGeneratorTool:
    def __init__(self):
        pass

    def artifact(
        self,
        query: str,
        sandbox_files: Optional[List[str]] = None,
        old_code: Optional[str] = None,
    ) -> Dict:
        """Generate a code artifact based on the provided input.

        Args:
            query (str): A description of the application you want to build.
            sandbox_files (Optional[List[str]], optional): A list of sandbox file paths. Defaults to None. Include these files if the code requires them.
            old_code (Optional[str], optional): The existing code to be modified. Defaults to None.

        Returns:
            Dict: A dictionary containing information about the generated artifact.
        """

        if old_code:
            user_message = f"{query}\n\nThe existing code is: \n```\n{old_code}\n```"
        else:
            user_message = query
        if sandbox_files:
            user_message += f"\n\nThe provided files are: \n{str(sandbox_files)}"

        messages: List[ChatMessage] = [
            ChatMessage(role="system", content=CODE_GENERATION_PROMPT),
            ChatMessage(role="user", content=user_message),
        ]
        try:
            sllm = Settings.llm.as_structured_llm(output_cls=CodeArtifact)  # type: ignore
            response = sllm.chat(messages)
            data: CodeArtifact = response.raw
            data_dict = data.model_dump()
            if sandbox_files:
                data_dict["files"] = sandbox_files
            return data_dict
        except Exception as e:
            logger.error(f"Failed to generate artifact: {str(e)}")
            raise e


def get_tools(**kwargs):
    return [FunctionTool.from_defaults(fn=CodeGeneratorTool().artifact)]
